# coding=utf-8
##
## Author: jmdvirus@aliyun.com
##
## Create: Thu Jun 17 19:37:31 2021
##

import keras
import tensorflow as tf
from keras.datasets import mnist
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation
from tensorflow.keras.optimizers import RMSprop

batch_size = 128
num_classes = 10
epochs = 20
img_size = 28*28

(x_train, y_train), (x_test, y_test) = mnist.load_data()
valid_len = 5000

print(x_train.shape)

x_len = x_train.shape[0]
train_len = x_len - valid_len
x_valid = x_train[train_len:]
y_valid = y_train[train_len:]

x_train = x_train[:train_len]
y_train = y_train[:train_len]

x_train = x_train.reshape(x_train.shape[0], img_size)
x_valid = x_valid.reshape(x_valid.shape[0], img_size)
x_test = x_test.reshape(x_test.shape[0], img_size)

x_train = x_train.astype('float32')
x_valid = x_valid.astype('float32')
x_test = x_test.astype('float32')

x_train /= 255
x_valid /= 255
x_test /= 255

y_train = tf.keras.utils.to_categorical(y_train, num_classes)
y_valid = tf.keras.utils.to_categorical(y_valid, num_classes)
y_test = tf.keras.utils.to_categorical(y_test, num_classes)

print("x_train.shape: {}".format(x_train.shape))
print("y_train.shape: {}".format(y_train.shape))
print("x_valid.shape: {}".format(x_valid.shape))
print("y_valid.shape: {}".format(y_valid.shape))
print("x_test.shape: {}".format(x_test.shape))
print("y_test.shape: {}".format(y_test.shape))

model = Sequential()
model.add(Dense(512, activation='relu', input_shape=(img_size,)))
model.add(Dropout(0.2))
model.add(Dense(512, activation='relu'))
model.add(Dropout(0.2))
model.add(Dense(num_classes, activation='softmax'))
model.summary()

model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
model.fit(x_train, y_train, epochs=epochs, batch_size=batch_size, verbose=1, validation_data=(x_valid, y_valid))

model.save("out/simple.rkbn")


